from deploy_script.conn.connection_pool import ConnectionPool
from deploy_script.modules.config import SparkConf
from deploy_script.logger_handler import Logger


def gen_key(configList):
    for config in configList:
        conn = ConnectionPool.get_conn(config.host)
        conn.run("rm -rf /root/.ssh")
        res = conn.run("ssh-keygen  -t rsa -P '' -f /root/.ssh/id_rsa")
        conn.run("cat ~/.ssh/id_rsa.pub >> ~/.ssh/authorized_keys")
        conn.run("chmod 600 ~/.ssh/authorized_keys")
        if res.ok:
            Logger.debug("{}生成密钥成功".format(config.host))


def gen_keys(configList):
    gen_key(configList)
    master = filter(lambda config: config.is_master, configList)
    master = list(master)
    if len(master) is not 1:
        raise Exception("master只能设置一个")
    slave = filter(lambda config: config.is_slave, configList)
    slave = list(slave)
    for m in master:
        master_conn = ConnectionPool.get_conn(m.host)
        master_conn.get("/root/.ssh/id_rsa.pub", 'id_rsa.pub')
        for i in slave:
            slave_conn = ConnectionPool.get_conn(i.host)
            slave_conn.put('id_rsa.pub', remote="/root/.ssh/id_rsa.pub")
            Logger.debug("主节点scp公钥到%s", i.host)

    for s in slave:
        conn = ConnectionPool.get_conn(s.host)
        conn.run("cat  ~/.ssh/id_rsa.pub >>  ~/.ssh/authorized_keys")
        conn.run("chmod 600 ~/.ssh/authorized_keys")
        Logger.debug("追加公钥")


if __name__ == '__main__':
    sparkConfigList = SparkConf.get_spark_conf_list()
    gen_keys(sparkConfigList)
